﻿#include "SubPixelEdgeExtract.h"
#include "opencv2/highgui.hpp"
#include <iostream>

namespace SubPixelEdgeExtract
{

static double xTable[8] = {cos(  0.0 * CV_PI / 180.0) * 0.5,
                    cos( 45.0 * CV_PI / 180.0) * sqrt(2.0) * 0.5,
                    cos( 90.0 * CV_PI / 180.0) * 0.5,
                    cos(135.0 * CV_PI / 180.0) * sqrt(2.0) * 0.5,
                    cos(180.0 * CV_PI / 180.0) * 0.5,
                    cos(225.0 * CV_PI / 180.0) * sqrt(2.0) * 0.5,
                    cos(270.0 * CV_PI / 180.0) * 0.5,
                    cos(315.0 * CV_PI / 180.0) * sqrt(2.0) * 0.5};

static double yTable[8] = {sin(  0.0 * CV_PI / 180.0) * 0.5,
                    sin( 45.0 * CV_PI / 180.0) * sqrt(2.0) * 0.5,
                    sin( 90.0 * CV_PI / 180.0) * 0.5,
                    sin(135.0 * CV_PI / 180.0) * sqrt(2.0) * 0.5,
                    sin(180.0 * CV_PI / 180.0) * 0.5,
                    sin(225.0 * CV_PI / 180.0) * sqrt(2.0) * 0.5,
                    sin(270.0 * CV_PI / 180.0) * 0.5,
                    sin(315.0 * CV_PI / 180.0) * sqrt(2.0) * 0.5};


static bool color2Gray(const cv::Mat &inImage, cv::Mat &gray)
{
    switch (inImage.channels())
    {
    case 1:
        gray = inImage;
        return true;
    case 3:
        cv::cvtColor(inImage, gray, cv::COLOR_RGB2GRAY);
        return true;
    case 4:
        cv::cvtColor(inImage, gray, cv::COLOR_RGBA2GRAY);
        return true;
    default:
        break;
    }
    return false;
}

bool edge(const cv::Mat &inImage, cv::Mat &edge8U, int threshold)
{
    if(inImage.empty()) return false;
    cv::Mat grayImage;
    if( !color2Gray(inImage, grayImage) ) return false;
    cv::Mat amp32F, angle32F;
    amplitudeAndAngle(grayImage, amp32F, angle32F);
//    cv::Mat amp;
//    amp32F.convertTo(amp, CV_8UC1, 1, 0);
//    cv::imshow("amp", amp);

//    cv::Mat ang;
//    angle32F.convertTo(ang, CV_8UC1, 0.5, 0);
//    cv::imshow("angle", ang);
    ampAngle2Edge(amp32F, angle32F, edge8U, threshold);

    return true;
}

void ampAngle2Edge(const cv::Mat &amp32F, const cv::Mat &angle32F, cv::Mat &edge8U, int threshold)
{
    edge8U.create(amp32F.rows, amp32F.cols, CV_8UC1);
    cv::Mat angle8U(angle32F.rows, angle32F.cols, CV_8UC1, cv::Scalar(0));
    angle45(amp32F, angle32F, angle8U, threshold);

    //cv::imshow("angle", angle8U * 31);
    int count=0;
    for (int r = 1; r < amp32F.rows - 1; r++)
    {
        // 3 * 3 邻域, 所以用3个指针, 一个指针指一行
        const float *pAmp1 = amp32F.ptr<float>(r - 1);
        const float *pAmp2 = amp32F.ptr<float>(r);
        const float *pAmp3 = amp32F.ptr<float>(r + 1);
        const uchar *pAng = angle8U.ptr<uchar>(r);
        uchar *pEdge = edge8U.ptr<uchar>(r);

        for (int c = 1; c < amp32F.cols - 1; c++)
        {
            pEdge[c] = 0;
            if(pAmp2[c] > threshold)
            {
                count ++;
                switch (pAng[c])
                {
                case ANGLE_270: // 270.
                    if (pAmp2[c] > pAmp1[c] && pAmp2[c] >= pAmp3[c])
                    {
                        pEdge[c] = 255;
                    }
                    break;
                case ANGLE_90: //90.
                    if (pAmp2[c] >= pAmp1[c] && pAmp2[c] > pAmp3[c])
                    {
                        pEdge[c] = 255;
                    }
                    break;
                case ANGLE_315: // 315.
                    if (pAmp2[c] > pAmp1[c + 1] && pAmp2[c] >= pAmp3[c - 1])
                    {
                        pEdge[c] = 255;
                    }
                    break;
                case ANGLE_135: //135.
                    if (pAmp2[c] >= pAmp1[c + 1] && pAmp2[c] > pAmp3[c - 1])
                    {
                        pEdge[c] = 255;
                    }
                    break;
                case ANGLE_0: //0.
                    if (pAmp2[c] >= pAmp2[c - 1] && pAmp2[c] > pAmp2[c + 1])
                    {
                        pEdge[c] = 255;
                    }
                    break;
                case ANGLE_180: //180.
                    if (pAmp2[c] > pAmp2[c - 1] && pAmp2[c] >= pAmp2[c + 1])
                    {
                        pEdge[c] = 255;
                    }
                    break;
                case ANGLE_45: //45.
                    if (pAmp2[c] >= pAmp1[c - 1] && pAmp2[c] > pAmp3[c + 1])
                    {
                        pEdge[c] = 255;
                    }
                    break;
                case ANGLE_225: //225.
                    if (pAmp2[c] > pAmp1[c - 1] && pAmp2[c] >= pAmp3[c + 1])
                    {
                        pEdge[c] = 255;
                    }
                    break;
                default:
                    break;
                }
            }
        }
    }

    std::cout<< "size = " << amp32F.cols * amp32F.rows << std::endl;
    std::cout << "count = " << count << std::endl;
    std::cout << "threshold = " << threshold << std::endl;
}

bool amplitudeAndAngle(const cv::Mat &inImage, cv::Mat &amp32F, cv::Mat &angle32F)
{
    if( inImage.empty() ) return false;
    cv::Mat grayImage;
    if( !color2Gray(inImage, grayImage) ) return false;
    cv::Mat sobelX, sobelY;
    cv::Sobel(grayImage, sobelX, CV_32F, 1, 0);
    cv::Sobel(grayImage, sobelY, CV_32F, 0, 1);
    cv::cartToPolar(sobelX, sobelY, amp32F, angle32F, true);
    return true;
}

bool subPixelEdge(const cv::Mat &inImage, cv::Mat &edge8U, cv::Mat &coordinate, int threshold)
{
    if(inImage.empty()) return false;
    cv::Mat grayImage;
    if( !color2Gray(inImage, grayImage) ) return false;
    cv::Mat amp32F, angle32F;
    amplitudeAndAngle(grayImage, amp32F, angle32F);
    ampAngle2SubPixelEdge(amp32F, angle32F, edge8U, coordinate, threshold);
    return true;
}

void angle45(const cv::Mat &amp32F, const cv::Mat &angle32F, cv::Mat &angle8U, int threshold)
{
    angle8U.create(angle32F.rows, angle32F.cols, CV_8UC1);// 以 45°为单位的角度
    for(int r = 0; r < angle32F.rows; r ++)
    {
        const float *pAmp = amp32F.ptr<float>(r);
        uchar *pAng2 = angle8U.ptr<uchar>(r);
        const float *pAng1 = angle32F.ptr<float>(r);
        for(int c = 0; c < angle32F.cols; c++)
        {
            if(pAmp[c] > threshold)
            {
                pAng2[c] = static_cast<uchar>( (22.5 + pAng1[c]) / 45.0);
                if(pAng2[c] > 8) pAng2[c] -= 8;
            }
            else
            {
                pAng2[c] = 0;
            }
        }
    }
}

void ampAngle2SubPixelEdge(const cv::Mat &amp32F, const cv::Mat &angle32F, cv::Mat &edge8U, std::vector<cv::Point2f> & points, int threshold)
{
    cv::Mat angle8U(angle32F.rows, angle32F.cols, CV_8UC1);// 以 45°为单位的角度
    edge8U.create(amp32F.rows, amp32F.cols, CV_8UC1);
    angle45(angle32F, angle32F, angle8U, threshold);
    for (int r = 1; r < amp32F.rows - 1; r++)
    {
        // 3 * 3 邻域, 所以用3个指针, 一个指针指一行
        const float *pAmp1 = amp32F.ptr<float>(r - 1);
        const float *pAmp2 = amp32F.ptr<float>(r);
        const float *pAmp3 = amp32F.ptr<float>(r + 1);
        const uchar *pAng = angle8U.ptr<uchar>(r);
        uchar *pEdge = edge8U.ptr<uchar>(r);
        for (int c = 1; c < amp32F.cols - 1; c++)
        {
            if(pAmp2[c] < threshold) continue;
            double dTmp = 0;
            switch (pAng[c])
            {
            case ANGLE_270://270.
                if (pAmp2[c] > pAmp1[c] && pAmp2[c] >= pAmp3[c])
                {
                    pEdge[c] = 255;
                    dTmp = -((double)pAmp1[c] - pAmp3[c]) / (pAmp1[c] + pAmp3[c] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_90://90.
                if (pAmp2[c] >= pAmp1[c] && pAmp2[c] > pAmp3[c])
                {
                    pEdge[c] = 255;
                    dTmp = ((double)pAmp1[c] - pAmp3[c]) / (pAmp1[c] + pAmp3[c] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_315://315.
                if (pAmp2[c] > pAmp1[c + 1] && pAmp2[c] >= pAmp3[c - 1])
                {
                    pEdge[c] = 255;
                    dTmp = -((double)pAmp1[c + 1] - pAmp3[c - 1]) / (pAmp1[c + 1] + pAmp3[c - 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_135://135.
                if (pAmp2[c] >= pAmp1[c + 1] && pAmp2[c] > pAmp3[c - 1])
                {
                    pEdge[c] = 255;
                    dTmp = ((double)pAmp1[c + 1] - pAmp3[c - 1]) / (pAmp1[c + 1] + pAmp3[c - 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_0://0.
                if (pAmp2[c] >= pAmp2[c - 1] && pAmp2[c] > pAmp2[c + 1])
                {
                    pEdge[c] = 255;
                    dTmp = ((double)pAmp2[c - 1] - pAmp2[c + 1]) / (pAmp2[c - 1] + pAmp2[c + 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_180://180.
                if (pAmp2[c] > pAmp2[c - 1] && pAmp2[c] >= pAmp2[c + 1])
                {
                    pEdge[c] = 255;
                    dTmp = -((double)pAmp2[c - 1] - pAmp2[c + 1]) / (pAmp2[c - 1] + pAmp2[c + 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_45://45.
                if (pAmp2[c] >= pAmp1[c - 1] && pAmp2[c] > pAmp3[c + 1])
                {
                    pEdge[c] = 255;
                    dTmp = ((double)pAmp1[c - 1] - pAmp3[c + 1]) / (pAmp1[c - 1] + pAmp3[c + 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_225://225.
                if (pAmp2[c] > pAmp1[c - 1] && pAmp2[c] >= pAmp3[c + 1])
                {
                    pEdge[c] = 255;
                    dTmp = -((double)pAmp1[c - 1] - pAmp3[c + 1]) / (pAmp1[c - 1] + pAmp3[c + 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            default:
                continue;
            }
            const double x = c + dTmp * xTable[pAng[c]];
            const double y = r + dTmp * yTable[pAng[c]];

            points.push_back(cv::Point(x, y));
        }
    }
}

void ampAngle2SubPixelEdge(const cv::Mat &amp32F, const cv::Mat &angle32F, cv::Mat &edge8U, cv::Mat &coordinate32F, int threshold)
{
    cv::Mat angle8U(angle32F.rows, angle32F.cols, CV_8UC1);// 以 45°为单位的角度
    edge8U.create(amp32F.rows, amp32F.cols, CV_8UC1);
    coordinate32F.create(amp32F.rows, amp32F.cols, CV_32FC2);
    angle45(amp32F, angle32F, angle8U, threshold);
    for (int r = 1; r < amp32F.rows - 1; r++)
    {
        // 3 * 3 邻域, 所以用3个指针, 一个指针指一行
        const float *pAmp1 = amp32F.ptr<float>(r - 1);
        const float *pAmp2 = amp32F.ptr<float>(r);
        const float *pAmp3 = amp32F.ptr<float>(r + 1);
        const uchar *pAng = angle8U.ptr<uchar>(r);
        uchar *pEdge = edge8U.ptr<uchar>(r);
        float *pCoordinate = coordinate32F.ptr<float>(r);
        for (int c = 1; c < amp32F.cols - 1; c++)
        {
            if(pAmp2[c] < threshold) continue;
            double dTmp = 0;
            switch (pAng[c])
            {
            case ANGLE_270://270.
                if (pAmp2[c] > pAmp1[c] && pAmp2[c] >= pAmp3[c])
                {
                    pEdge[c] = 255;
                    dTmp = -((double)pAmp1[c] - pAmp3[c]) / (pAmp1[c] + pAmp3[c] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_90://90.
                if (pAmp2[c] >= pAmp1[c] && pAmp2[c] > pAmp3[c])
                {
                    pEdge[c] = 255;
                    dTmp = ((double)pAmp1[c] - pAmp3[c]) / (pAmp1[c] + pAmp3[c] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_315://315.
                if (pAmp2[c] > pAmp1[c + 1] && pAmp2[c] >= pAmp3[c - 1])
                {
                    pEdge[c] = 255;
                    dTmp = -((double)pAmp1[c + 1] - pAmp3[c - 1]) / (pAmp1[c + 1] + pAmp3[c - 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_135://135.
                if (pAmp2[c] >= pAmp1[c + 1] && pAmp2[c] > pAmp3[c - 1])
                {
                    pEdge[c] = 255;
                    dTmp = ((double)pAmp1[c + 1] - pAmp3[c - 1]) / (pAmp1[c + 1] + pAmp3[c - 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_0://0.
                if (pAmp2[c] >= pAmp2[c - 1] && pAmp2[c] > pAmp2[c + 1])
                {
                    pEdge[c] = 255;
                    dTmp = ((double)pAmp2[c - 1] - pAmp2[c + 1]) / (pAmp2[c - 1] + pAmp2[c + 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_180://180.
                if (pAmp2[c] > pAmp2[c - 1] && pAmp2[c] >= pAmp2[c + 1])
                {
                    pEdge[c] = 255;
                    dTmp = -((double)pAmp2[c - 1] - pAmp2[c + 1]) / (pAmp2[c - 1] + pAmp2[c + 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_45://45.
                if (pAmp2[c] >= pAmp1[c - 1] && pAmp2[c] > pAmp3[c + 1])
                {
                    pEdge[c] = 255;
                    dTmp = ((double)pAmp1[c - 1] - pAmp3[c + 1]) / (pAmp1[c - 1] + pAmp3[c + 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            case ANGLE_225://225
                if (pAmp2[c] > pAmp1[c - 1] && pAmp2[c] >= pAmp3[c + 1])
                {
                    pEdge[c] = 255;
                    dTmp = -((double)pAmp1[c - 1] - pAmp3[c + 1]) / (pAmp1[c - 1] + pAmp3[c + 1] - 2 * pAmp2[c]);
                    break;
                }
                else
                {
                    continue;
                }
            default:
                continue;
            }
            const double x = c + dTmp * xTable[pAng[c]];
            const double y = r + dTmp * yTable[pAng[c]];
            pCoordinate[2 * c] = (float)x;
            pCoordinate[2 * c + 1] = (float)y;
        }
    }
}

}
